package com.lynn.xss.filter;

import org.apache.commons.text.StringEscapeUtils;
import org.springframework.util.StringUtils;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.util.HashMap;
import java.util.Map;

/**
 * @ClassName XssHttpServletRequestWrapper
 * @Description 重写wrapper
 * @Author ly
 * @Date 2019-05-29 11:00
 * @Version 1.0
 */
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {
    /**
     * Constructs a request object wrapping the given request.
     *
     * @param request The request to wrap
     * @throws IllegalArgumentException if the request is null
     */
    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
    }

    @Override
    public String[] getParameterValues(String parameter) {
        String[] values = super.getParameterValues(parameter);
        if (values == null) {
            return null;
        }
        int count = values.length;
        String[] encodedValues = new String[count];
        for (int i = 0; i < count; i++) {
            encodedValues[i] = StringEscapeUtils.unescapeHtml4(values[i]);
        }
        return encodedValues;

    }

    @Override
    public String getParameter(String name) {
        String oldValue = super.getParameter(name);
        if (!StringUtils.isEmpty(oldValue)) {
            oldValue = StringEscapeUtils.unescapeHtml4(oldValue);
        }
        return oldValue;
    }

    @Override
    public String getHeader(String name) {
        String value = super.getHeader(name);
        if (value == null) {
            return null;
        }
        return StringEscapeUtils.unescapeHtml4(value);
    }

    @Override
    public Map getParameterMap() {
        HashMap paramMap = (HashMap) super.getParameterMap();
        paramMap = (HashMap) paramMap.clone();
        for (Object o : paramMap.entrySet()) {
            Map.Entry entry = (Map.Entry) o;
            String[] values = (String[]) entry.getValue();
            int length = values.length;
            for (int i = 0; i < length; i++) {
                if (values[i] != null) {
                    values[i] = StringEscapeUtils.unescapeHtml4(values[i]);
                }
            }
            entry.setValue(values);
        }
        return paramMap;
    }
}
